import torch
import torch.nn as nn
from orbitP.script import config
from orbitP.layers.FeatureEncoder import FeatureEncoder


class MLP(nn.Module):
    """
    纯 MLP 基线，支持异步输入/输出。
    - 输入: x_enc [B, T_enc, D_in]
    - 输出: y_hat [B, T_dec, output_dim]
    注意: x_dec/x_mark_dec 仅用于接口对齐（可传入），模型本身不使用它们。
    """
    def __init__(self, args):
        super(MLP, self).__init__()
        # 保留与现有代码一致的超参读取方式
        self.input_window = config.training_length            # 历史窗口长度（T_enc）
        self.output_window = config.predicting_length  # 预测窗口长度（T_dec）
        self.input_dim = args.head_dim               # 每个时间步的特征维度（D_in）
        self.output_dim = config.outputSize + config.nll  # 每步预测维度
        self.hidden_dim = args.d_model
        # 用于与其它模型一致的前处理（如果 FeatureEncoder 改变特征）
        self.FeatureEncoder = FeatureEncoder(args)

        # 展平后输入维度
        in_dim = self.input_window * self.input_dim
        out_dim = self.output_window * self.output_dim

        layers = []
        # 输入层
        layers.append(nn.Linear(in_dim, args.d_model))
        layers.append(nn.ReLU())
        layers.append(nn.Dropout(args.dropout))

        # 隐藏层
        for _ in range(args.num_layers - 2):
            layers.append(nn.Linear(args.d_model, args.d_model))
            layers.append(nn.ReLU())
            layers.append(nn.Dropout(args.dropout))

        # 输出层
        layers.append(nn.Linear(args.d_model, out_dim))
        self.mlp = nn.Sequential(*layers)

    def forward(self, x_enc, x_mark_enc, x_dec=None, x_mark_dec=None):
        input = self.FeatureEncoder(x_enc, x_mark_enc, x_dec, x_mark_dec)  # [B, T_enc, D_in]
        B, L, N = input.shape  # B L N
        # 如果 FeatureEncoder 改变了时间长度或维度，请确保与 config.seq_len / args.head_dim 对齐
        # 展平历史窗口
        input = input.reshape(B, -1)  # [B, T_enc * D_in]

        # 直接用 MLP 预测整个未来窗口
        output = self.mlp(input)       # [B, T_dec * output_dim]

        # 恢复时间维度
        output = output.view(B, self.output_window, self.output_dim)  # [B, T_dec, output_dim]
        return output